from src.utils import *
from src.layers import *
import matplotlib.pyplot as plt
import numpy as np
# choose TIP model: 'cat' - TIP-cat
#					'add' - TIP-add
MOD = 'cat'
MAX_EPOCH = 180 # 180

# set training device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
AUPRC = []
AUROC = []
LRlist = []
bestAUPRC = 0
bestPRCLR = 10.0
bestAUROC = 0
bestROCLR = 10.0
Dd = 0
print("bestAUPRC = AUPRC:{} AUPRC Dd:{} AUROC:{} AUROC Dd:{}".format(AUPRC, Dd, AUROC, Dd))
LRcont = 0
if LRcont == 0:
#for Dd in range(1, 21):
	#break
	#LRcont = 200
	#LR = (LRcont*1.0/1000)
	#LR = 0.5
	#print(Dd, end="\n")

	# initial model
	# best lr = 0.01
	if MOD == 'cat':
		settings = Setting(sp_rate=0.9, lr=0.01, prot_drug_dim=16, n_embed=48, n_hid1=32, n_hid2=16, num_base=32)
		model = TIP(settings, device) #best Dd = 16
	else:
		settings = Setting(sp_rate=0.9, lr=0.01, prot_drug_dim=64, n_embed=64, n_hid1=32, n_hid2=16, num_base=32)
		model = TIP(settings, device, mod='add') #best Dd = 64


	#print("total parameters number is {}".format(sum(x.numel() for x in model.parameters())))

	# initial optimizer
	optimizer = torch.optim.Adam(model.parameters(), lr=settings.lr)
	min_loss = 100
	times = 100
	time = 0
	# train TIP model
	'''
	for e in range(MAX_EPOCH):
		model.train()
		optimizer.zero_grad()
		loss = model()
		loss.backward()
		optimizer.step()
	'''
	for e in range(MAX_EPOCH):
		model.train()
		optimizer.zero_grad()
		loss = model()
		if min_loss < loss.item() :
			time += 1
		else :
			time  = 0
			min_loss = loss.item()
		if time >= times:
			break
		else:
			pass
		print(loss.item())
		loss.backward()
		optimizer.step()

	#evaluate on test set
	'''
		bestAUPRC = 0
		bestPRCLR = 10.0
		bestAUROC = 0
		bestROCLR = 10.0
	'''
	record, auprc, auroc, ap = model.test()
	if auprc > bestAUPRC:
		bestAUPRC = auprc
		bestPRCLR = Dd
	if auroc > bestAUROC:
		bestAUROC = auroc
		bestROCLR = Dd
	AUPRC.append(auprc)

	AUROC.append(auroc)
	LRlist.append(Dd)


#print("bestAUPRC = AUPRC:{} AUPRC Dd:{} AUROC:{} AUROC Dd:{}".format(bestAUPRC, Dd, bestAUROC, Dd))
AUPRC = np.array(AUPRC)
AUROC = np.array(AUROC)
LRlist = np.array(LRlist)
plt.plot(LRlist, AUPRC, color="red", linewidth=2)
plt.plot(LRlist, AUROC, color="blue", linewidth=2)
plt.figure(figsize = (20, 12))  # 定义图的大小
plt.ylabel("red:AUPRC blue:AUPOC")
plt.xlabel("Learning Rate")
plt.title("Dd - AUPRC/AUROC")
plt.plot(LRlist, AUPRC)
plt.plot(LRlist, AUROC)
plt.savefig("compare.png")
# print((model.data.dd_test_range).shape)
# print((model.test_neg_index))
# save trained model
torch.save(model, f'./saved_model/tip-{model.mod}-example.pt')
